package SimpleLACore

import chisel3._
import chisel3.util._
import chisel3.util.random.LFSR

class SimpleCacheAXI(val id: Int) extends Module with SimpleLACoreParam {
  val io = IO(new Bundle {
    val core = new SimpleCacheIO
    val mem = new AXI4Bundle
  })
  require((nSets & (nSets - 1)) == 0, "nSets must be power of 2")
  require((lineWords & (lineWords - 1)) == 0, "wordsPerLine must be power of 2")
  require((nWays & (nWays - 1)) == 0, "nWays must be power of 2")

  def getIdx(addr: UInt):UInt = addr(idxBits + lgLineBytes - 1, lgLineBytes)
  def getTag(addr: UInt):UInt = addr(31, 32 - tagBits)
  def getOffset(addr: UInt):UInt = addr(lgLineBytes - 1, 2)

  val rand = LFSR(8)(lgnWays - 1, 0)

  val index = WireInit(UInt(idxBits.W), DontCare)
  val metas = Seq.fill(nWays)(Mem(nSets, new SimpleCacheMeta)(index))
  val datum = Seq.fill(nWays)(Mem(nSets, Vec(lineWords, Vec(4, UInt(8.W))))(index))
  val reset_counter = RegInit(0.U((idxBits + 1).W))
  val reset_done = reset_counter(idxBits)

  val coreBusyReg = RegInit(false.B)
  val coreReqReg = RegEnable(io.core.req.bits, io.core.req.valid)
  val coreBusy = coreBusyReg || io.core.req.valid
  val coreReq = Mux(io.core.req.valid, io.core.req.bits, coreReqReg)

  val buf = Reg(Vec(lineWords, Vec(4, UInt(8.W))))
  val bufBase = Reg(UInt(lgLineWords.W))
  val bufAddr = Reg(UInt((32 - lgLineBytes).W))
  val bufHead = Reg(Vec(lineWords, Bool()))
  val bufHeadNext = (bufHead.reverse.tail :+ bufHead.last).reverse
  val bufValid = RegInit(VecInit.tabulate(lineWords)(_=>false.B))
  val bufDirty = Reg(Bool())

  val sReady :: sFill :: sSwap :: sWb :: Nil = Enum(4)
  val state = RegInit(sReady)
  val doResp = RegInit(false.B)
  val aFire = RegInit(false.B)
  val wFire = RegInit(false.B)

  io.mem.ar.valid := false.B
  io.mem.aw.valid := false.B
  io.mem.w.valid  := false.B
  io.mem.ar.bits := DontCare
  io.mem.aw.bits := DontCare
  io.mem.w.bits  := DontCare
  io.mem.aw.bits.id := id.U
  io.mem.aw.bits.size := 2.U
  io.mem.aw.bits.lock := 0.U
  io.mem.aw.bits.cache := 0.U
  io.mem.aw.bits.prot := 0.U
  io.mem.ar.bits.id := id.U
  io.mem.ar.bits.size := 2.U
  io.mem.ar.bits.lock := 0.U
  io.mem.ar.bits.cache := 0.U
  io.mem.ar.bits.prot := 0.U
  io.mem.r.ready  := true.B
  io.mem.b.ready  := true.B
  io.core.resp.valid := false.B
  io.core.resp.bits := DontCare

  when(io.core.resp.valid) {
    coreBusyReg := false.B
  }.elsewhen(io.core.req.valid) {
    coreBusyReg := true.B
  }
  //reset
  when(!reset_done){
    reset_counter := reset_counter + 1.U
    index := reset_counter
    metas.foreach(_.valid := false.B)
  }
  //req
  when(reset_done && coreBusy){
    when(coreReq.cacop) {
      when(state === sReady){//Only do CACOP when ready
        io.core.resp.valid := true.B
        switch(coreReq.code) {
          is(0.U) {
            index := getIdx(coreReq.addr)
            (metas zip UIntToOH(coreReq.addr(lgnWays - 1, 0)).asBools).foreach{ case (meta, sel) =>
              when(sel){meta.valid := false.B}
            }
          }
          is(1.U) {
            index := getIdx(coreReq.addr)
            val waySel = UIntToOH(coreReq.addr(lgnWays - 1, 0)).asBools
            (metas zip waySel).foreach{ case (meta, hit) =>
              when(hit){meta.valid := false.B}
            }
            val meta = Mux1H(waySel, metas)
            val data = Mux1H(waySel, datum)

            when(meta.dirty && meta.valid) {
              state := sWb
              buf := data//.asTypeOf(buf)
              bufAddr := Cat(meta.tag, index)
              bufBase := getOffset(coreReq.addr)
              bufHead := UIntToOH(getOffset(coreReq.addr)).asBools
              (bufValid zip UIntToOH(getOffset(coreReq.addr)).asBools).foreach( x => x._1 := ! x._2)
            }
          }
          is(2.U) {
            index := getIdx(coreReq.addr)
            val (dirty, data) = (datum zip metas).foldLeft((false.B, 0.U((lineWords * 32).W))) { case ((dirtyIn, dataIn), (line, meta)) =>
              val hit = meta.valid && meta.tag === getTag(coreReq.addr)
              when(hit){meta.valid := false.B}
              (meta.dirty & hit | dirtyIn, line.asUInt & Fill(lineWords * 32, hit) | dataIn)
            }
            when(dirty){
              state := sWb
              buf := data.asTypeOf(buf)
              bufAddr := coreReq.addr(31, lgLineBytes)
              bufBase := getOffset(coreReq.addr)
              bufHead := UIntToOH(getOffset(coreReq.addr)).asBools
              (bufValid zip UIntToOH(getOffset(coreReq.addr)).asBools).foreach( x => x._1 := ! x._2)
            }
          }
        }
      }
    }.elsewhen(coreReq.mat =/= 0.U){//CC
      when(!(state === sFill && bufValid.reduceTree(_&_))){
        index := getIdx(coreReq.addr)
        //buffer hit?
        val bufHit = bufAddr === coreReq.addr(31, lgLineBytes) && bufValid(getOffset(coreReq.addr))
        val bufData = buf(getOffset(coreReq.addr)).asUInt & Fill(32, bufHit)
        when(bufHit){
          (buf(getOffset(coreReq.addr)) zip toBytes(coreReq.wdata) zip coreReq.wen.asBools).foreach{
            case ((target, wData), wen) => when(wen){target := wData}
          }
          bufDirty := coreReq.wen.orR || bufDirty
        }
        //cache hit?
        val (hit, data) = (datum zip metas).foldLeft((bufHit, bufData)) { case ((hitIn, dataIn), (line, meta)) =>
          val hit = meta.valid && meta.tag === getTag(coreReq.addr)
          val data = line(getOffset(coreReq.addr))
          when(hit){
            (line(getOffset(coreReq.addr)) zip toBytes(coreReq.wdata) zip coreReq.wen.asBools).foreach{
              case ((target, wData), wen) => when(wen){target := wData}
            }
            meta.dirty := coreReq.wen.orR || meta.dirty
          }
          (hit | hitIn, data.asUInt & Fill(32, hit) | dataIn)
        }
        when(hit){
          io.core.resp.valid := true.B
          io.core.resp.bits := data
        }.elsewhen(state === sReady){
          state := sFill
          bufAddr := coreReq.addr(31, lgLineBytes)
          bufBase := getOffset(coreReq.addr)
          bufHead := UIntToOH(getOffset(coreReq.addr)).asBools
          bufDirty := false.B
          when(coreReq.preld) {io.core.resp.valid := true.B}
        }
      }
    }.elsewhen(state === sReady){//SUC
      when(!aFire){
        when(coreReq.wen.orR()){
          io.mem.aw.valid := true.B
          io.mem.aw.bits.addr := coreReq.addr
          io.mem.aw.bits.len := 0.U
          io.mem.aw.bits.burst := 1.U
        }.otherwise{
          io.mem.ar.valid := true.B
          io.mem.ar.bits.addr := coreReq.addr
          io.mem.ar.bits.len := 0.U
          io.mem.ar.bits.burst := 1.U
        }
      }
      when(!wFire && coreReq.wen.orR()){
        io.mem.w.valid := true.B
        io.mem.w.bits.data := coreReq.wdata
        io.mem.w.bits.strb := coreReq.wen
        io.mem.w.bits.last := true.B
      }
      when(io.mem.r.valid){
        aFire := false.B
        io.core.resp.valid := true.B
        io.core.resp.bits := io.mem.r.bits.data
      }.elsewhen(io.mem.aw.fire || io.mem.ar.fire){
        aFire := true.B
      }
      when(io.mem.b.valid){
        wFire := false.B
        aFire := false.B
        io.core.resp.valid := true.B
      }.elsewhen(io.mem.w.fire){
        wFire := true.B
      }
    }
  }

  when(reset_done){
    switch(state) {
      is(sFill) {
        when(!aFire){
          io.mem.ar.valid := true.B
          io.mem.ar.bits.addr := Cat(bufAddr, bufBase, 0.U(2.W))
          io.mem.ar.bits.len := (lineWords - 1).U
          io.mem.ar.bits.burst := 2.U
          when(io.mem.ar.ready){
            aFire := true.B
          }
        }
        when(io.mem.r.valid) {
          (buf zip bufValid zip bufHead).foreach{ case ((word, valid), isHead) =>
            when(isHead){
              word := io.mem.r.bits.data.asTypeOf(word)
              valid := true.B
            }
          }
          bufHead := bufHeadNext
        }
        //swap buffer and cache line (sSwap)
        when(bufValid.reduceTree(_&_)){
          val replaceWay = rand // random replacement
          index := bufAddr // scala will truncate tag bits
          val waySel = UIntToOH(replaceWay).asBools
          //write data and meta
          (datum zip waySel).foreach{ case (line, sel) =>
            when(sel){
              line := buf
            }
          }
          (metas zip waySel).foreach{ case (meta, sel) =>
            when(sel){
              meta.tag := bufAddr(tagBits + idxBits - 1, idxBits)
              meta.dirty := bufDirty
              meta.valid := true.B
            }
          }
          //prepare buffer
          val meta = Mux1H(waySel, metas)
          buf := Mux1H(waySel, datum)
          bufAddr := Cat(meta.tag, index)

          aFire := false.B
          when(meta.dirty && meta.valid) {
            state := sWb
            (bufValid zip bufHead).foreach{ case (valid, isHead) => when(isHead){valid := false.B}}
          }.otherwise {
            state := sReady
            bufValid.foreach(_ := false.B)
          }
        }
      }
      is(sWb) {
        when(!aFire){
          io.mem.aw.valid := true.B
          io.mem.aw.bits.addr := Cat(bufAddr, bufBase, 0.U(2.W))
          io.mem.aw.bits.len := (lineWords - 1).U
          io.mem.aw.bits.burst := 2.U
          when(io.mem.aw.ready){
            aFire := true.B
          }
        }
        when(!wFire) {
          io.mem.w.valid := true.B
          io.mem.w.bits.data := Mux1H(bufHead, buf).asUInt
          io.mem.w.bits.last := !bufValid.reduceTree(_|_)
          io.mem.w.bits.strb := 0xf.U
          when(io.mem.w.ready) {
            bufHead := bufHeadNext
            (bufValid zip bufHeadNext).foreach{ case (valid, isHead) => when(isHead){valid := false.B}}
            when(!bufValid.reduceTree(_|_)) {
              wFire := true.B
            }
          }
        }
        when(io.mem.b.valid){
          state := sReady
          aFire := false.B
          wFire := false.B
        }
      }
    }
  }
}